from ultralytics import YOLO

# 加载模型
model_path = 'data/weights/best.pt'
model = YOLO(model_path)

# 导出为 TorchScript，指定设备和尺寸
output_path = 'data/weights/best_converted.torchscript'
model.export(
    format='torchscript',
    imgsz=640,         # 输入图像尺寸
    device='cpu',      # 强制使用 CPU
    optimize=True      # 启用优化
)

print(f"Model exported to {output_path}")

# 验证
import torch
converted_model = torch.jit.load(output_path)
print("Converted model loaded successfully")
print(type(converted_model))